import torch
import torch.nn as nn 
import torch.nn.functional as F 
import argparse, os, random
import torch.optim as optimizer 
from sklearn.metrics import accuracy_score
from torch.utils.data import DataLoader,Subset, dataloader,random_split
import numpy as np 

import sys
sys.path.append("core/data_loader")

from multimodal.model.t4sa.core.data_loader.load_data_t4sa import *
from core.model.UniModel import Model_Single_Text
from core.model.basic_multi_modal import Multi_Fusion_Model_Baseline
from core.model.image import * 
from core.model.decent_model_gate_cofeature import * 

def parse_args():
    parser = argparse.ArgumentParser()
    # Model
    parser.add_argument('--layer', type=int, default=4)
    parser.add_argument('--hidden_size', type=int, default=1024)
    parser.add_argument('--dim', type=int, default=1024)

    parser.add_argument('--dropout_r', type=float, default=0.1)
    parser.add_argument('--multi_head', type=int, default=8)
    parser.add_argument('--ff_size', type=int, default=2048)
    parser.add_argument('--word_embed_size', type=int, default=300)

    parser.add_argument('--image_hw', type=int, default=256)
    parser.add_argument('--patch_hw', type=int, default=32)
    parser.add_argument('--lang_size', type=int, default=31)
    parser.add_argument('--num_classes', type=int, default=3)
    parser.add_argument('--img_len', type=int, default=65)
    parser.add_argument('--text_len', type=int, default=31)

    # Training
    parser.add_argument('--output', type=str, default='ckpt/')
    parser.add_argument('--name', type=str, default='exp_cofeature/')
    parser.add_argument('--batch_size', type=int, default=64)
    parser.add_argument('--max_epoch', type=int, default=99)
    parser.add_argument('--opt', type=str, default="Adam")
    parser.add_argument('--opt_params', type=str, default="{'betas': '(0.9, 0.98)', 'eps': '1e-9'}")
    parser.add_argument('--lr_base', type=float, default=0.0001)
    parser.add_argument('--lr_decay', type=float, default=0.5)
    parser.add_argument('--lr_decay_times', type=int, default=2)
    parser.add_argument('--warmup_epoch', type=float, default=0)
    parser.add_argument('--grad_norm_clip', type=float, default=-1)
    parser.add_argument('--eval_start', type=int, default=0)
    parser.add_argument('--early_stop', type=int, default=3)
    parser.add_argument('--seed', type=int, default=random.randint(0, 9999999))
    parser.add_argument('--sigma', type=float, default=1.0)
    parser.add_argument('--ans_size', type=float, default=3)
    parser.add_argument('--pred_func',type=str, default="amax")
    # Dataset and task
    parser.add_argument('--root_dir',type=str,default="/mnt/ssd/Datasets_/")
    args = parser.parse_args()
    return args


def set_seed(seed):
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)  
    np.random.seed(seed)  
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True 





def evaluate_single(net, eval_loader, args,device):
    accuracy = []
    net.train(False)
    preds = []
    for step, (
            x,
            ans,
    ) in enumerate(eval_loader):
        x = x.to(device)
        pred = net(x,False).cpu().data.numpy()

        
        ans = ans.cpu().data.numpy()
        accuracy += list(np.argmax(pred, axis=1) == ans)

   
    return 100*np.mean(np.array(accuracy)) #
    

if __name__ == "__main__":
    args = parse_args()

    val_dataset =T4SA_Only_Vision(1,None)
    test_dataset =T4SA_Only_Vision(2,None)



    print("train_data_len:",val_dataset.__len__())
    print("val_data_len:",test_dataset.__len__())

   
    valid_data_iter = DataLoader(val_dataset,batch_size=128,shuffle=False,num_workers=4) #
    test_data_iter = DataLoader(test_dataset,batch_size=128,shuffle=False)
    
    
    device = torch.device('cuda:1')

    net = Model_Single_Vision(args).to(device)

    net.load_state_dict(torch.load("ckpt/exp0/best99.89880245044576_719174.pkl")['state_dict'])

    print(evaluate_single(net, valid_data_iter, args,device))